# %%
import queue
import threading
import angr
from angrutils import *
import numpy
import networkx
import conf
from bidict import bidict

# %%

class MyInstruction():
    # https://cs.brown.edu/courses/cs033/docs/guides/x64_cheatsheet.pdf
    # https://cch123.gitbooks.io/duplicate/content/part3/translation-details/function-calling-sequence/xmm-registers.html
    s_regs = ['al', 'bl', 'cl', 'dl', 'sil', 'dil', 'spl', 'bpl',
              'r8b', 'r9b', 'r10b', 'r11b', 'r12b', 'r13b', 'r14b', 'r15b',
              'ah', 'bh', 'ch', 'dh',
              'ax', 'bx', 'cx', 'dx', 'si', 'di', 'sp', 'bp',
              'r8w', 'r9w', 'r10w', 'r11w', 'r12w', 'r13w', 'r14w', 'r15w',
              'eax', 'ebx', 'ecx', 'edx', 'esi', 'edi', 'esp', 'ebp',
              'r8d', 'r9d', 'r10d', 'r11d', 'r12d', 'r13d', 'r14d', 'r15d',
              'rax', 'rbx', 'rcx', 'rdx', 'rsi', 'rdi', 'rsp', 'rbp',
              'r8', 'r9', 'r10', 'r11', 'r12', 'r13', 'r14', 'r15',
              'xmm0', 'xmm1', 'xmm2', 'xmm3', 'xmm4', 'xmm5', 'xmm6', 'xmm7',
              'xmm8', 'xmm9', 'xmm10', 'xmm11', 'xmm12', 'xmm13', 'xmm14', 'xmm15']

    def __init__(self, ins=None):
        self.mnemonic = ""
        self.ops = []
        if ins is not None:
            self.parse(ins)

    def __changeReg(self, op: str) -> str:
        if op in MyInstruction.s_regs:
            return "REG"
        return None

    def __changeImm(self, op: str) -> int:
        ret = None
        try:
            ret = int(op, 16)
        except:
            try:
                ret = int(op, 10)
            except:
                try:
                    ret = int(op, 8)
                except:
                    try:
                        ret = int(op, 2)
                    except:
                        print(op, 'is not Imm')
                        ret = None

        return ret

    def __changeMem(self, op: str) -> str:
        if '[' in op and ']' in op:
            return 'MEM'
        return None

    def __norm(self, op: str) -> str:
        '''
        op一共有四种情况：
        1. Imm 立即数
        2. E 寄存器
        3. R[E] 存储在寄存器E中的值
        4. M[x] 存储在地址x的地址
        '''
        ret = self.__changeReg(op)
        if ret is None:
            ret = self.__changeMem(op)
            if ret is None:
                ret = self.__changeImm(op)
                if ret is None:
                    print("'" + op + "'" + " is not REG IMM MEM")
                else:
                    ret = str(ret)
        return ret

    def parse(self, capstoneinsn: angr.block.CapstoneInsn):
        self.mnemonic = capstoneinsn.mnemonic
        self.ops = [self.__norm(i)
                    for i in capstoneinsn.op_str.split(", ") if i != '']


    def compare(self, other) -> int:
        '''
        比较两条指令，计算分数
        Instead of comparing original instructions, we choose to
        compare the normalized instructions. The first advantage
        is more resistance to register reassignment, which is very
        common in compiler optimization. Second, we want to do a
        fuzzy matching. This is different from what David and Ya-
        hav did in [12], where they use exact matching when com-
        paring operands. Besides, we allow partial matching. For
        example, we give a score of 5 to instruction pair cmp [eax],0
        and cmp ebx, 0, although they are two types of instructions.
        The first instruction is comparing an immediate value with
        a memory reference while the second with an register. The
        reason for allowing partial matching is that, even for the
        same variable, compilers have the freedom to represent it as
        a register variable or a memory variable. Allowing partial
        matching can tolerate these differences.
        '''
        score = 0

        if self.mnemonic == other.mnemonic:
            score = score + conf.SCORE_IDENTICAL_MNEMONIC
            for op1, op2 in zip(self.ops, other.ops):
                if op1 == op2:
                    if type(op1) == int:
                        score = score + conf.SCORE_IDENTICAL_CONSTANT
                    else:
                        score = score + conf.SCORE_IDENTICAL_OPERAND

        return score

# %%


class MyBlock():
    def __init__(self, block: angr.block.Block) -> None:
        self.instructions = []
        disassemble_code = block.disassembly

        for ins in disassemble_code.insns:
            self.__add(ins)

        self.compare_self_score = self.compare(self)

    def __add(self, ins):
        self.instructions.append(MyInstruction(ins))

    def compare(self, other) -> int:
        i = 1
        j = 1
        M = numpy.zeros((len(self.instructions) + 1,
                        len(other.instructions) + 1))
        for i in range(1, len(self.instructions) + 1):
            # j_start = j
            j_start = 1  # TODO 究竟是按照指令顺序进行判断还是乱序指令是个问题
            for j in range(j_start, len(other.instructions) + 1):
                M[i][j] = max(self.instructions[i-1].compare(
                    other.instructions[j-1]) + M[i-1][j-1],
                    M[i-1][j],
                    M[i][j-1])

        return M[len(self.instructions)][len(other.instructions)]

# %%


class MyFunction():
    def __init__(self, func_node: angr.knowledge_plugins.functions.function.Function):
        self.func_node = func_node
        self.longest_path = self.__dfs()  # 存储的是node节点
        self.longest_path_my_block = [MyBlock(self.func_node.get_block(b.addr)) for b in self.longest_path]
        self.num_instructions = sum([len(b.instructions) for b in self.longest_path_my_block])
        self.len_blocks = len(list(self.func_node.blocks))
        self.longest_path_score = self.__path_score(self.longest_path, self.func_node)

        _scoress = [MyBlock(self.func_node.get_block(b.addr)).compare_self_score for b in func_node.blocks]
        self.all_self_score = sum(_scoress) # 计算函数自身的分数
        # if self.longest_path_score == 0:
        #     print(func_node.addr, "longest_path_score is 0")
        self.lcs_score = 0
        self.lcs_score_path = []
        self.other_func_node = None
        # print(self.longest_path_len())

    def __remove(self):
        '''
        删除环路。为啥要用dfs？因为dfs可以避免将构成最长路径的边给删除掉。
        这样处理之后最终只能拿到一个相对最长的边，因为无法避免将构成最长路径的边删除的可能。
        '''
        dcg = self.func_node.graph
        while True:
            try:
                c = networkx.find_cycle(dcg)
                dcg.remove_edge(c[-1][0], c[-1][1])
            except:
                break

    def __dfs(self):
        if networkx.algorithms.is_directed_acyclic_graph(self.func_node.graph) == False:
            self.__remove()
        return networkx.algorithms.dag_longest_path(self.func_node.graph)
    
    def longest_path_len(self):
        return len(self.longest_path)

    def SameDegree(self, u, v, other_func_node) -> bool:
        g1 = self.func_node.graph
        g2 = other_func_node.graph
        ind1 = g1.in_degree(v)
        ind2 = g2.in_degree(u)
        outd1 = g1.out_degree(v)
        outd2 = g2.out_degree(u)
        return (abs(ind1 - ind2) == 0) and (abs(outd1 - outd2) == 0)

        # return (len(u.predecessors()) == len(v.predecessors())) and\
        #     (len(u.successors()) == len(v.successors()))

    def __LCS(self, other_func_node, u, delta, came_from):

        for i in range(1, len(self.longest_path) + 1):
            v = self.longest_path[i - 1]
            sim = 0
            if self.SameDegree(u, v, other_func_node):
                sim = MyBlock(other_func_node.get_block(u.addr)).compare(
                    MyBlock(self.func_node.get_block(v.addr)))
            delta[u][i] = max(
                delta[came_from[u]][i - 1] + sim,
                delta[came_from[u]][i],
                delta[u][i - 1]
            )

    def __path_score(self, p, func_node)->int:
        score = 0
        for n in p:
            block_node = MyBlock(func_node.get_block(n.addr))
            score = score + block_node.compare(block_node)
        # score = 0
        # for n in func_node.blocks:
        #     b = MyBlock(n)
        #     score = score + b.compare(b)
        return score

    def normScore(self)->float:
        score1 = self.__path_score(self.lcs_score_path, self.other_func_node)
        return self.lcs_score / (score1 + self.longest_path_score - self.lcs_score)

    def pathExploration(self, other_my_func_node):
        self.lcs_score = 0
        self.lcs_score_path = []
        self.other_func_node = other_my_func_node.func_node

        G = other_my_func_node.func_node.graph
        first_node = other_my_func_node.func_node.get_node(other_my_func_node.func_node.addr)

        came_from = {}
        came_from[first_node] = None

        delta = {}
        delta[None] = numpy.zeros((len(self.longest_path) + 1))

        sigma = {}
        Q = queue.Queue()
        Q.put(first_node)
        while Q.empty() == False:
            u = Q.get()
            if u not in delta:
                delta[u] = numpy.zeros((len(self.longest_path) + 1))
            self.__LCS(other_my_func_node.func_node, u, delta, came_from)

            try:
                if sigma[u] < delta[u][len(self.longest_path)]:
                    sigma[u] = delta[u][len(self.longest_path)]
                    for succ in G.successors(u):
                        Q.put(succ)
                        came_from[succ] = u # 修改当前节点的父节点
            except:
                sigma[u] = delta[u][len(self.longest_path)]


        tmp = sorted(sigma.items(), key=lambda item: item[1])
        self.lcs_score = tmp[-1][1]
        self.lcs_score_path = []

        u = tmp[-1][0]
        n = u
        while n is not None:
            self.lcs_score_path.append(n)
            n = came_from[n]
        
        return self.lcs_score


#%%

class MyCFG():
    def __init__(self, bin_path):
        self.proj = angr.Project(bin_path, load_options={'auto_load_libs': False})
        self.cfg = self.proj.analyses.CFGFast(normalize=True)

        print("This is the graph:", self.cfg.graph)
        print("It has %d nodes and %d edges" %
            (len(self.cfg.graph.nodes()), len(self.cfg.graph.edges())))

        # this grabs *any* node at a given location:
        entry_node = self.cfg.model.get_any_node(self.proj.entry)

        # on the other hand, this grabs all of the nodes
        print("There were %d contexts for the entry block" %
            len(self.cfg.model.get_all_nodes(self.proj.entry)))

        # we can also look up predecessors and successors
        print("Predecessors of the entry point:", entry_node.predecessors)
        print("Successors of the entry point:", entry_node.successors)
        print("Successors (and type of jump) of the entry point:", [jumpkind + " to " + str(
            node.addr) for node, jumpkind in self.cfg.model.get_successors_and_jumpkind(entry_node)])
        
        self.func_names = []
        self.name2my_functions = {}
        self.my_functions = []
        self.my_functions2functions = bidict()
        for func_node in self.cfg.functions.values():
            # if func_node.name.startswith("__"):
            #     continue
            # else:
            my_func_node = MyFunction(func_node)
            if my_func_node.num_instructions < 3: # 过滤掉指令小于3的函数
                continue
            self.my_functions.append(my_func_node)
            self.my_functions2functions[my_func_node] = func_node
            self.name2my_functions[func_node.name] = my_func_node
            self.func_names.append(func_node.name)
        print("函数节点数量：", len(self.my_functions))


        

    def get_by_name(self, name: str):
        try:
            return self.name2my_functions[name]
        except:
            return None



#%%

def filt(my_function_node, other_my_functions):
    ret_my_function_nodes = []

    score = my_function_node.all_self_score
    for i in other_my_functions:
        _ratio =  min(score / i.all_self_score, i.all_self_score / score)
        if _ratio > 0.5:
            if _ratio < 0.75 and abs(i.len_blocks - my_function_node.len_blocks) > 5:
                continue
            ret_my_function_nodes.append(i)

    return ret_my_function_nodes

#%%
def compare(my_function_node, other_my_functions):
    norm_scores = {}
    for j in other_my_functions:
        my_function_node.pathExploration(j)
        norm_scores[j.func_node.name] = my_function_node.normScore()
    return norm_scores


#%%

def compareCFG(cfg1, cfg2, NUM_THREADS = 2):

    norm_scores_q = queue.Queue()

    def task2(my_func_node, other_my_functions):
        filt_function_nodes = other_my_functions
        filt_function_nodes = filt(my_func_node, cfg2.my_functions)
        if len(filt_function_nodes) == 0:
            print(my_func_node.func_node.name, "过滤之后没有匹配函数")
        return compare(my_func_node, filt_function_nodes)

    def task1(my_func_nodes):
        norm_scores = {}
        for it,i in enumerate(my_func_nodes):
            norm_scores[i.func_node.name] = task2(i, cfg2.my_functions)
            if it % 100 == 0:
                print("已经处理", it, "个节点")
        norm_scores_q.put(norm_scores)

    q = queue.Queue()
    i = 0
    STEP = len(cfg1.my_functions) // NUM_THREADS + 1
    while i < len(cfg1.my_functions):
        q.put(cfg1.my_functions[i:i+STEP])
        i = i + STEP

    threads = []
    while q.empty() == False:
        datas = q.get()
        threads.append(threading.Thread(target=task1, args=(datas,)))
        threads[-1].start()

    print("等待线程结束")
    for t in threads:
        t.join()

    norm_scores = {}
    while norm_scores_q.empty() == False:
        n = norm_scores_q.get()
        norm_scores.update(n)

    print("所有节点计算完毕")
    return norm_scores

#%%
def top10scores(norm_scores):
    norm_scores_sorted = sorted(norm_scores.items(), key=lambda item: item[1], reverse=True)
    max_score = norm_scores_sorted[0]
    ret = norm_scores_sorted[:10]
    for i in range(11, len(norm_scores_sorted)):
        if norm_scores_sorted[i][1] > (max_score[1] * 0.98):
            ret.append(norm_scores_sorted[i])
    return ret

#%%
def my_print(norm_scores, my_cfg, other_my_cfg):
    ret_matchs = []
    num_dismatch = 0
    for i in my_cfg.func_names:
        _tmp = sorted(norm_scores[i].items(), key=lambda item: item[1])
        if len(_tmp) == 0:
            print("no score for ", i)
            print("\n\n\n")
            num_dismatch = num_dismatch + 1
            continue
        _k_v = _tmp[-1]
        _ret = top10scores(norm_scores[i])
        _top10_names = [_i[0] for _i in _ret]
        if i not in _top10_names:
            num_dismatch = num_dismatch + 1

            _i_node = my_cfg.get_by_name(i)
            _k_v_node = other_my_cfg.get_by_name(_k_v[0])

            print('dismatch:', i, _k_v[0])
            print('dismatch:', _i_node.func_node.addr, _k_v_node.func_node.addr, "score:", _k_v[1])
            
            
            try:
                print('match score:', norm_scores[i][i])
            except:
                if other_my_cfg.get_by_name(i) is not None:
                    print(i,"被过滤掉了")
                else:
                    print(i, ': no such function in', other_my_cfg.proj.filename)
                
            print(_top10_names)
            
            print("\n\n")
        else:
            ret_matchs.append(i)

    ratio = num_dismatch / len(my_cfg.my_functions)
    print("dismatch ratio:", ratio)
    return (ratio, ret_matchs)


#%%

def print_func_node(func_node):
    print(func_node.name)
    for b in func_node.blocks:
        b.pp()
        print("\n\n")
    func_node.dbg_draw('images/test.png')

